#!/usr/bin/env python

import sys

from matplotlib import pyplot as plt
from mpl_toolkits.basemap import Basemap
import numpy as np

def on_help(p_name):
    print "Usage:", p_name, "setting_file, station_file, out_file"

def parse_settings(settings_filename):
    "Returns 2 dicts with basemap setting and domain_settings"
    map_settings = {}
    domain_settings = {}
    f = open(settings_filename)
    for l in f:
        if l.startswith('#') or l.strip() == '':
            continue
        key, raw_val=l.rstrip().split('=')
        try:
            true_val = float(raw_val)
        except ValueError:
            true_val = raw_val
        if key.startswith('domain'):
            domain_settings[key] = true_val
        else:
            map_settings[key] = true_val
    f.close()
    return map_settings, domain_settings

def parse_stations(station_filename):
    "Returns a list of tuples of stations lon and lat"
    stations_lon_lat = []
    f = open(station_filename)
    for l in f:
        if l.startswith('#') or l.strip() == '':
            continue
        lon_s, lat_s = l.rstrip().split()
        stations_lon_lat.append((float(lon_s), float(lat_s)))
    return stations_lon_lat

if len(sys.argv) < 4:
    on_help(sys.argv[0])
    sys.exit(1)
p_name, setting_filename, station_filename, out_filename = sys.argv


map_settings, domain_settings = parse_settings(setting_filename)    
print map_settings, domain_settings

stations_lon_lat = parse_stations(station_filename)
print stations_lon_lat

x_offset = (map_settings['width'] - domain_settings['domain_res'] * domain_settings['domain_x_num']) / 2.0
y_offset = (map_settings['height'] - domain_settings['domain_res'] * domain_settings['domain_y_num']) / 2.0

domain_grids_x, domain_grids_y = np.meshgrid(np.arange(domain_settings['domain_x_num']), np.arange(domain_settings['domain_y_num']))
domain_grids_x *= domain_settings['domain_res']
domain_grids_y *= domain_settings['domain_res']
domain_grids_x += x_offset
domain_grids_y += y_offset

domain_grids_x_sparse = domain_grids_x[::10, ::10]
domain_grids_y_sparse = domain_grids_y[::10, ::10] 


fig = plt.figure(figsize=(24,18))

#If you have map data of shapefile format, try it like following
#m.readshapefile('./map/bou2_4p','bou2_4p')
#m.readshapefile('./map/diquJie_polyline','diquJie_polyline',color='#007F7F')

m = Basemap(projection='lcc', **map_settings)
m.drawcoastlines()
m.bluemarble()

# plot the stations
stations_on_map = [m(lon, lat) for lon, lat in stations_lon_lat]
m_x, m_y = zip(*stations_on_map) 
m.plot(m_x, m_y, 'ro')

# plot the grids
m.plot(domain_grids_x, domain_grids_y,'w+')
m.plot(domain_grids_x_sparse, domain_grids_y_sparse,'+', color='orange')

plt.savefig(out_filename)
plt.close()
sys.exit(0)
